import torchvision
from torch import nn
from torch.utils.tensorboard import SummaryWriter

from common.logging_format import *
from model_ import *

logger = Logger("train_log.log", "train_log")
train_dataset = torchvision.datasets.CIFAR10(root='../data/CIFAR10', train=True, download=True,
                                             transform=torchvision.transforms.ToTensor())

test_dataset = torchvision.datasets.CIFAR10(root='../data/CIFAR10', train=False, download=True,
                                            transform=torchvision.transforms.ToTensor())
logger.info("训练集数据长度:{}".format(len(train_dataset)))
logger.info("测试集数据长度:{}".format(len(test_dataset)))

train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=64)
test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=64)

my_net = MyNet()
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(my_net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)

epochs = 5

total_test_step = 0

writer = SummaryWriter("../logs")

total_train_step = 0
for epoch in range(epochs):
    logger.info('--------------------第{}轮训练--------------------'.format(epoch + 1))

    for images, labels in train_dataloader:
        output = my_net(images)
        loss = loss_fn(output, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if total_train_step % 100 == 0:
            logger.info(
                'Epoch [{}/{}], batch [{:4d}] Loss: {:.4f}'.format(epoch + 1, epochs, total_train_step, loss.item()))
            writer.add_scalar('train_loss', loss.item(), total_train_step)
        total_train_step += 1
    total_test_loss = 0
    total_accuracy = 0
    with torch.no_grad():
        for data, target in test_dataloader:
            output = my_net(data)
            loss = loss_fn(output, target)
            total_test_loss += loss.item()
            # 计算正确率
            accuracy = (output.argmax(1) == target).sum()
            total_accuracy += accuracy
    logger.info('Test Loss: {:.4f}'.format(total_test_loss / len(test_dataset)))
    logger.info('Total accuracy: {:.4f}'.format(total_accuracy / len(test_dataset)))
    writer.add_scalar('test_loss', total_test_loss, total_test_step)
    writer.add_scalar('test_accuracy', total_accuracy, total_test_step)
    total_test_step += 1

    torch.save(my_net, 'model_{}.pth'.format(epoch))
writer.close()
